/*
* Copyright 1993-2010 NVIDIA Corporation.  All rights reserved.
*
* NOTICE TO USER:
*
* This source code is subject to NVIDIA ownership rights under U.S. and
* international Copyright laws.  Users and possessors of this source code
* are hereby granted a nonexclusive, royalty-free license to use this code
* in individual and commercial software.
*
* NVIDIA MAKES NO REPRESENTATION ABOUT THE SUITABILITY OF THIS SOURCE
* CODE FOR ANY PURPOSE.  IT IS PROVIDED "AS IS" WITHOUT EXPRESS OR
* IMPLIED WARRANTY OF ANY KIND.  NVIDIA DISCLAIMS ALL WARRANTIES WITH
* REGARD TO THIS SOURCE CODE, INCLUDING ALL IMPLIED WARRANTIES OF
* MERCHANTABILITY, NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
* IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL,
* OR CONSEQUENTIAL DAMAGES, OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
* OF USE, DATA OR PROFITS,  WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
* OR OTHER TORTIOUS ACTION,  ARISING OUT OF OR IN CONNECTION WITH THE USE
* OR PERFORMANCE OF THIS SOURCE CODE.
*
* U.S. Government End Users.   This source code is a "commercial item" as
* that term is defined at  48 C.F.R. 2.101 (OCT 1995), consisting  of
* "commercial computer  software"  and "commercial computer software
* documentation" as such terms are  used in 48 C.F.R. 12.212 (SEPT 1995)
* and is provided to the U.S. Government only as a commercial end item.
* Consistent with 48 C.F.R.12.212 and 48 C.F.R. 227.7202-1 through
* 227.7202-4 (JUNE 1995), all U.S. Government End Users acquire the
* source code with only those rights set forth herein.
*
* Any use of this source code in individual and commercial software must
* include, in the user documentation and internal comments to the code,
* the above Disclaimer and U.S. Government End Users Notice.
*/

// GPU Debug API Test for CUDA
// Based on the matrixMulDrv sample from the CUDA SDK


// Change to control which device this program is run on.
#define CUDA_DEVICE_ORDINAL 0

#include <cuda.h>

#include "matrixMul.h"
#include "matrixMul_general.h"

// includes, std
#include <math.h>
#include <crtdbg.h>

#include <map>
#include <string>
#include <iostream>

#include "../../../Common/check.h"
#define ALIGN_UP(offset, alignment) \
    (offset) = ( (offset) + (alignment) - 1 ) & ~( (alignment) - 1 )

////////////////////////////////////////////////////////////////////////////////

extern "C"
void computeGold(float*, const float*, const float*, unsigned int, unsigned int, unsigned int);

// Initializes a matrix with random float entries.
void randomInit(float* data, int size)
{
    for (int i = 0; i < size; ++i)
    {
        data[i] = rand() / (float)RAND_MAX;
    }
}

//  Convenient function for waiting for keystroke.
void PressAKeyToContinue()
{
    std::wcout << L"<<< press enter to exit >>>" << std::endl;
    getchar();
}


////////////////////////////////////////////////////////////////////////////////
//  Launch helper function
////////////////////////////////////////////////////////////////////////////////

bool LaunchMatrixMul(CUfunction func)
{
    unsigned int size_A = WA * HA;
    unsigned int mem_size_A = sizeof(float) * size_A;
    unsigned int size_B = WB * HB;
    unsigned int mem_size_B = sizeof(float) * size_B;
    unsigned int size_C = WC * HC;
    unsigned int mem_size_C = sizeof(float) * size_C;

    float* h_A = 0;
    float* h_B = 0;
    float* h_C = 0;
    CUdeviceptr d_A = 0;
    CUdeviceptr d_B = 0;
    CUdeviceptr d_C = 0;
    CUresult status = CUDA_ERROR_UNKNOWN;

    // set seed for rand()
    srand(2006);

    //  Breakable block.
    for (;;)
    {
        // allocate host memory for matrices A and B
        h_A = (float*) malloc(mem_size_A);
        CheckConditionBreak(h_A);
        h_B = (float*) malloc(mem_size_B);
        CheckConditionBreak(h_B);

        // initialize host memory
        randomInit(h_A, size_A);
        randomInit(h_B, size_B);

        // allocate device memory
        status = cuMemAlloc( &d_A, mem_size_A );
        CheckConditionBreak(CUDA_SUCCESS == status);
        cuMemAlloc( &d_B, mem_size_B );
        CheckConditionBreak(CUDA_SUCCESS == status);

        // copy host memory to device
        status = cuMemcpyHtoD( d_A, h_A, mem_size_A );
        CheckConditionBreak(CUDA_SUCCESS == status);
        status = cuMemcpyHtoD( d_B, h_B, mem_size_B );
        CheckConditionBreak(CUDA_SUCCESS == status);

        // allocate device memory for result
        h_C = (float*) malloc(mem_size_C);
        CheckConditionBreak(h_C);

        // clear output buffer
        status = cuMemAlloc(&d_C, mem_size_C);
        CheckConditionBreak(CUDA_SUCCESS == status);
        ::memset(h_C, 0, mem_size_C);
        status = cuMemcpyHtoD(d_C, h_C, mem_size_C);
        CheckConditionBreak(CUDA_SUCCESS == status);

        // setup execution parameters
        status = cuFuncSetBlockShape( func, BLOCK_SIZE, BLOCK_SIZE, 1 );
        CheckConditionBreak(CUDA_SUCCESS == status);
        status = cuFuncSetSharedSize( func, 2*BLOCK_SIZE*BLOCK_SIZE*sizeof(float) );
        CheckConditionBreak(CUDA_SUCCESS == status);

        //  TODO: check return codes.  These are pretty tame calls though.
        int offset = 0;
        void* ptr = 0;
        int tmp = 0;

        ptr = (void*)d_C;
        ALIGN_UP( offset, __alignof(ptr) );
        cuParamSetv( func, offset,  &d_C, sizeof(d_C) );
        offset += sizeof(d_C);

        ptr = (void*)d_A;
        ALIGN_UP( offset, __alignof(ptr) );
        cuParamSetv( func, offset,  &d_A, sizeof(d_A) );
        offset += sizeof(d_A);

        ptr = (void*)d_B;
        ALIGN_UP( offset, __alignof(ptr) );
        cuParamSetv( func, offset,  &d_B, sizeof(d_B) );
        offset += sizeof(d_B);

        tmp = WA;
        ALIGN_UP( offset, __alignof( tmp ) );
        cuParamSeti( func, offset, tmp );
        offset += sizeof(tmp);

        tmp = WB;
        ALIGN_UP( offset, __alignof( tmp ) );
        cuParamSeti( func, offset, tmp );
        offset += sizeof(tmp);

        cuParamSetSize( func, offset );

        std::wcout << L"Launching grid..." << std::endl;
        status = cuLaunchGrid( func, WC / BLOCK_SIZE, HC / BLOCK_SIZE );
        CheckConditionBreak(CUDA_SUCCESS == status);

        status = cuCtxSynchronize();
        CheckConditionBreak(CUDA_SUCCESS == status);
        std::wcout << L"...Grid Completed" << std::endl;

        status = cuMemcpyDtoH((void *) h_C, d_C, mem_size_C);
        CheckConditionBreak(CUDA_SUCCESS == status);

        status = cuCtxSynchronize();
        CheckConditionBreak(CUDA_SUCCESS == status);

        // compute reference solution
        float* reference = (float*) malloc(mem_size_C);
        computeGold(reference, h_A, h_B, HA, WA, WB);

        //  Cleanup.
        cuMemFree(d_C);
        cuMemFree(d_B);
        cuMemFree(d_A);
        free(reference);
        free(h_C);
        free(h_B);
        free(h_A);

        return true;
    }

    //TODO: cleanup code on failure!
    return false;
}


////////////////////////////////////////////////////////////////////////////////
// Program main
////////////////////////////////////////////////////////////////////////////////
bool main_0(const std::wstring& applicationDir, int argc, wchar_t** argv)
{
    CUresult status = CUDA_ERROR_UNKNOWN;
    bool br = false;

    //  Initialize CUDA.

    status = cuInit(0);
    CheckConditionXR(CUDA_SUCCESS == status, 0);

    CUdevice cuDevice = 0;
    CUcontext cuContext = 0;
    br = CreateCudaContext(CUDA_DEVICE_ORDINAL, &cuDevice, &cuContext);
    CheckConditionXR(br, false);

    CUmodule cuModuleMatrixMul = 0;

    //  From here, we have a breakable block to make sure we clean up.
    for (;;)
    {
        br = LoadMatrixMulModule(&cuModuleMatrixMul, cuDevice, applicationDir);
        CheckConditionBreak(br);

        CUfunction funcMatrixMul = 0;

        br = (cuModuleGetFunction(&funcMatrixMul, cuModuleMatrixMul, "matrixMul") == CUDA_SUCCESS);
        CheckConditionBreak(br);

        br = LaunchMatrixMul(funcMatrixMul);
        CheckConditionBreak(br);

        bool getUserInput = ShouldGetUserInput(argc, argv);

        if (getUserInput)
        {
            PressAKeyToContinue();
        }

        break;
    }

    if (cuModuleMatrixMul)
    {
        cuModuleUnload(cuModuleMatrixMul);
    }
    if (cuContext)
    {
        cuCtxDestroy(cuContext);
    }

    // return the results
    return br;
}

int wmain(int argc, wchar_t** argv)
{
    bool br = false;

    std::wstring applicationPath = argv[0];
    std::wstring applicationDir = applicationPath.substr(0, applicationPath.find_last_of(L"\\") + 1);

    br = main_0(applicationDir, argc, argv);

    exit((br) ? EXIT_SUCCESS : EXIT_FAILURE);
}
